"""便于直接通过字符串配置网络结构以及训练流程
"""
import transformers
import logging
import torch.nn as nn
import torch.optim as optim
from typing import Dict


class PluginManager:
    """插件管理，通过字典存储所有插件.
    """
    def __init__(self) -> None:
        # 预定义一些常用配置
        self.plugin_container: Dict[str: Dict[str: object]] = {
            "act": {
                "relu": nn.ReLU,
                "prelu": nn.PReLU,
                "tanh": nn.Tanh,
                "sigmoid": nn.Sigmoid,
            },
            "optimizer": {
                "AdamW": optim.AdamW,
                "Adam": optim.Adam,
                "SGD": optim.SGD,
            },
            "scheduler": {
                "constant_warmup": transformers.get_constant_schedule_with_warmup,
                "cosine_warmup": transformers.get_cosine_schedule_with_warmup,
                "linear_warmup": transformers.get_linear_schedule_with_warmup
            }
        }

    def register(self, plugin_type: str, plugin_name: str, cls):
        print("register", plugin_type, plugin_name, cls)
        if plugin_type not in self.plugin_container:
            self.plugin_container[plugin_type] = {}
        self.plugin_container[plugin_type][plugin_name] = cls
        logging.debug(self.plugin_container)
    
    def get(self, plugin_type: str, plugin_name: str):
        if plugin_type in self.plugin_container and plugin_name in self.plugin_container[plugin_type]:
            return self.plugin_container[plugin_type][plugin_name]
        else:
            logging.error(f"{plugin_type} and {plugin_name} not in plugins")
            return None


DefaultPluginManager = PluginManager()

def register_plugin(plugin_type: str, plugin_name: str):
    def decorator(cls):
        DefaultPluginManager.register(plugin_type, plugin_name, cls)
        return cls
    
    return decorator


def get_plugin(plugin_type: str, plugin_name: str):
    logging.debug(f"get plugin: {plugin_type}, {plugin_name}")

    return DefaultPluginManager.get(plugin_type, plugin_name)

